import os
import argparse
import matplotlib
import numpy as np
import sys
import shutil
import random
import math
from pathlib import Path
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from tqdm import tqdm
import utils.autoanchor as autoAC


def show_width_height(data, bins=50):
    """
    Display bouding box distribution with histgram.
    """
    if data.dtype != np.float32:
        data = data.astype(np.float32)
    width = data[:, 0]
    height = data[:, 1]
    ratio = width / height
    plt.figure(1, figsize=(24, 6))
    plt.hist(ratio, bins=bins, color='magenta')
    plt.xlabel('Width / Height ')
    plt.ylabel('number')
    plt.title('Distribution of aspect ratio(Width / Height )')
    plt.savefig('shape-distribution.png')
    plt.show()
    return ratio


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-filelist', default=
        '../data/jingshanchunji/202204/test/labels', help=
        """path to filelist
""")
    parser.add_argument('-num_clusters', default=9, type=int, help=
        'number of clusters\n')
    parser.add_argument('-size', default=1024, type=int, help='input size\n')
    args = parser.parse_args()
    wh_list = []
    most_ratio = 0.0
    ratio_dic = {}
    for root, dirs, files in os.walk(args.filelist):
        loop = tqdm(files)
        for file in loop:
            filename = os.path.join(root, file)
            f = open(filename)
            for line in f.readlines():
                line = line.rstrip('\n')
                w, h = line.split(' ')[3:]
                ratio = 0
                if float(h) != 0:
                    ratio = float(eval(w) / eval(h))
                    if ratio_dic.get(round(ratio), 0) == 0:
                        ratio_dic[round(ratio)] = 1
                    else:
                        ratio_dic[round(ratio)] = ratio_dic[round(ratio)] + 1
                    wh_list.append((eval(w), eval(h)))
            f.close()
            loop.set_postfix(ratio=ratio, width=w, height=h)
    ratio_dic = sorted(ratio_dic.items(), key=lambda x: -x[0])
    print(ratio_dic)
    most_ratio = ratio_dic[0][0]
    print(most_ratio)
    show_width_height(np.array(wh_list), bins=500)
    print('Start to do kmeans, please wait for a moment.')
    new_anchors = autoAC.kmean_anchors('../data/VOC.yaml', args.
        num_clusters, args.size, most_ratio, 10000)
    print('new_anchors:\n', new_anchors)


if __name__ == '__main__':
    main()
